/*
 * twoLvAnalyser.cc
 *
 *  Created on: Dec 6, 2012
 *      Author: hxin
 */

#include "twoLvHTAnalyser.h"
#include <iostream>
#include <cassert>
#include <new>

twoLvHTAnalyser::twoLvHTAnalyser() {
}

twoLvHTAnalyser::twoLvHTAnalyser(string refFileName, int shortHTLength,
		int maxError, int pruneThreshold) {
	this->longHTLength = shortHTLength * 3;
	this->midHTLength = shortHTLength * 2;
	this->shortHTLength = shortHTLength;

	this->maxError = maxError;
	this->pruneThreshold = pruneThreshold;

	this->longHT = refAnalyser(refFileName, longHTLength, 0);
	this->midHT = refAnalyser(refFileName, midHTLength, 0);
	this->shortHT = refAnalyser(refFileName, shortHTLength, 0);

	this->conductedAF.assign(maxError + 1, 0);
}

twoLvHTAnalyser::twoLvHTAnalyser(string refFileName, int shortHTLength,
		int midHTLength, int longHTLength, int maxError, int pruneThreshold) {
	this->longHTLength = longHTLength;
	this->midHTLength = midHTLength;
	this->shortHTLength = shortHTLength;

	this->maxError = maxError;
	this->pruneThreshold = pruneThreshold;

	this->longHT = refAnalyser(refFileName, longHTLength, 0);
	this->midHT = refAnalyser(refFileName, midHTLength, 0);
	this->shortHT = refAnalyser(refFileName, shortHTLength, 0);

	this->conductedAF.assign(maxError + 1, 0);
}

twoLvHTAnalyser::twoLvHTAnalyser(const twoLvHTAnalyser& rhs) {
	this->longHTLength = rhs.longHTLength;
	this->midHTLength = rhs.midHTLength;
	this->shortHTLength = rhs.shortHTLength;

	this->maxError = rhs.maxError;
	this->pruneThreshold = rhs.pruneThreshold;

	this->longHT = rhs.longHT;
	this->midHT = rhs.midHT;
	this->shortHT = rhs.shortHT;

	this->conductedAF = rhs.conductedAF;
}

twoLvHTAnalyser::~twoLvHTAnalyser() {
}

twoLvHTAnalyser& twoLvHTAnalyser::operator =(const twoLvHTAnalyser & rhs) {
	this->longHTLength = rhs.longHTLength;
	this->midHTLength = rhs.midHTLength;
	this->shortHTLength = rhs.shortHTLength;

	this->maxError = rhs.maxError;
	this->pruneThreshold = rhs.pruneThreshold;

	this->longHT = rhs.longHT;
	this->midHT = rhs.midHT;
	this->shortHT = rhs.shortHT;

	this->conductedAF = rhs.conductedAF;
	return *this;
}

refHash* twoLvHTAnalyser::clone() {
	return new twoLvHTAnalyser(*this);
}

void twoLvHTAnalyser::loadHash(int chrNum) {
	//longHT.loadHash(chrNum);
	midHT.loadHash(chrNum);
	shortHT.loadHash(chrNum);
	pruneHT();
}

void twoLvHTAnalyser::unloadHash() {
	//longHT.unloadHash();
	midHT.unloadHash();
	shortHT.unloadHash();
}

vector<unsigned long long> twoLvHTAnalyser::query(string read) {
	return query(read, conductedAF);
}

vector<unsigned long long> twoLvHTAnalyser::query(string read,
		vector<unsigned long int>& conductedAF) {
	vector<unsigned long long> locList;

	int KmerNum = read.length() / shortHTLength;
	int* cumulatedAF;

	cumulatedAF = new int[maxError + 1];

	keyStruct * sortInput;

	try {
		sortInput = new keyStruct[KmerNum];
	} catch (const bad_alloc& ba) {
		cerr << "Caught not enough memory" << endl;
		cerr << "KmerNum = " << KmerNum << endl;
		cerr << "shortHTLength = " << shortHTLength << endl;
		cerr << "read.length() = " << read.length() << endl;
		cerr.flush();
		exit(1);
	}
	//Getting the Kmers
	for (int i = 0; i < KmerNum; i++) {
		string KmerString = read.substr(i * shortHTLength, shortHTLength);
		sortInput[i].key_entry_size = shortHT.query(KmerString).size();
		sortInput[i].original_order = i;
	}
	//Sort them
	qsort(sortInput, KmerNum, sizeof(keyStruct), compareEntrySize);

	for (int setError = 0; setError <= maxError; setError++) {
		cumulatedAF[setError] = 0;
		for (int e = 0; e <= setError; e++) {

			cumulatedAF[setError] += sortInput[e].key_entry_size;

			if (sortInput[e].key_entry_size > pruneThreshold) {
				for (int i = setError + 1; i < KmerNum; i++) {
					if (sortInput[i].original_order
							== sortInput[e].original_order + 1) {
						string KmerString = read.substr(
								sortInput[e].original_order * shortHTLength,
								midHTLength);
						cumulatedAF[setError] -= sortInput[e].key_entry_size;
						cumulatedAF[setError] += midHT.query(KmerString).size();
						break;
					} else if (sortInput[i].original_order
							== sortInput[e].original_order - 1) {
						string KmerString = read.substr(
								sortInput[i].original_order * shortHTLength,
								midHTLength);
						cumulatedAF[setError] -= sortInput[e].key_entry_size;
						cumulatedAF[setError] += midHT.query(KmerString).size();
						break;
					}
				}
			}
		}

		conductedAF[setError] += cumulatedAF[setError];
	}

	delete[] sortInput;
	try {
		locList.assign(cumulatedAF[maxError], 0);
	} catch (const bad_alloc& ba) {
		cerr << "cumulatedAF[maxError]: " << cumulatedAF[maxError] << endl;
	}
	delete[] cumulatedAF;
	return locList;
}

vector<unsigned long long> twoLvHTAnalyser::testLv(string Kmer) {
	string shortKmer = Kmer.substr(0, shortHTLength);
	string midKmer = Kmer.substr(0, midHTLength);
//	string longKmer = Kmer.substr(0, longHTLength);

	if (shortHT.query(shortKmer).size() <= pruneThreshold)
		return shortHT.query(shortKmer);
	else {
		assert(midHT.query(midKmer).size() > 0);
		return midHT.query(midKmer);
	}

}

void twoLvHTAnalyser::genStat() {
//longHT.genStat();
	midHT.genStat();
	shortHT.genStat();
}

void twoLvHTAnalyser::printStat() {
//cout << "Printing longHT stats:" << endl;
//longHT.printStat();
	cout << "Printing midHT stats:" << endl;
	midHT.printStat();
	cout << "Printing shortHT stats:" << endl;
	shortHT.printStat();
	cout << "***AF info:" << endl;
	for (int e = 0; e <= maxError; e++)
		cout << "For e = " << e << ", conducted AF after CKS: "
				<< conductedAF[e] << endl;
}

void twoLvHTAnalyser::pruneHT() {
	vector<string> longPruneSet;
	vector<string> midPruneSet;

	/*//Prune the long ht
	 longHT.begin();
	 while (!longHT.end()) {
	 string longKmer = longHT.getKmerString();
	 string shortKmer = longKmer.substr(0, shortHTLength);
	 string midKmer = longKmer.substr(0, midHTLength);
	 if (shortHT.query(shortKmer).size() < pruneThreshold) {
	 longPruneSet.push_back(longKmer);
	 midPruneSet.push_back(midKmer);
	 } else if (midHT.query(midKmer).size() < pruneThreshold) {
	 longPruneSet.push_back(midKmer);
	 }
	 longHT.next();
	 }

	 for (int i = 0; i < longPruneSet.size(); i++)
	 longHT.pruneKmer(longPruneSet[i]);
	 for (int i = 0; i < midPruneSet.size(); i++)
	 midHT.pruneKmer(midPruneSet[i]);
	 */

	midPruneSet.clear();
	midHT.begin();
	while (!midHT.end()) {
		string midKmer = midHT.getKmerString();
		string shortKmer = midKmer.substr(0, shortHTLength);
		if (shortHT.query(shortKmer).size() <= pruneThreshold)
			midPruneSet.push_back(midKmer);
		midHT.next();
	}
	for (int i = 0; i < midPruneSet.size(); i++)
		midHT.pruneKmer(midPruneSet[i]);
}
